package class01;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

public class Code04_UnionFind {

	// 样本进来会包一层，叫做元素。   相当于套个圈
	public static class Element<V> {
		public V value;

		public Element(V value) {
			this.value = value;
		}

	}


	// 图中的箭头可以用Map也可以用单链表

	public static class UnionFindSet<V> {
		public HashMap<V, Element<V>> elementMap; // a -> a圈
		// key：某个元素  value：该元素的父
		public HashMap<Element<V>, Element<V>> fatherMap;
		// key 某个集合的代表元素(就是根)  value：该集合的大小
		public HashMap<Element<V>, Integer> sizeMap;


		/**
		 * 并查集的初始化
		 * @param list
		 */
		public UnionFindSet(List<V> list) {
			elementMap = new HashMap<>();
			fatherMap = new HashMap<>();
			sizeMap = new HashMap<>();
			for (V value : list) {
				Element<V> element = new Element<V>(value);
				elementMap.put(value, element);
				fatherMap.put(element, element);
				sizeMap.put(element, 1);
			}
		}


		// 给定一个element，往上一直找，把代表元素返回
		// 这里面含有扁平优化 即路径压缩
		// 这里还可以使用递归写法
		private Element<V> findHead(Element<V> element) {
			Stack<Element<V>> path = new Stack<>();
			while (element != fatherMap.get(element)) { // 一直往上窜
				path.push(element);
				element = fatherMap.get(element);
			}

			// 扁平优化
			while (!path.isEmpty()) {
				fatherMap.put(path.pop(), element); // 将沿途的每个元素都直接连接到大顶部
			}
			return element;
		}

		/**
		 * 判断某两个元素是否在同一集合
		 * @param a
		 * @param b
		 * @return
		 */
		public boolean isSameSet(V a, V b) {
			if (elementMap.containsKey(a) && elementMap.containsKey(b)) { // 并查集能使用的前提是已经初始化
				return findHead(elementMap.get(a)) == findHead(elementMap.get(b));
			}
			return false;
		}

		/**
		 * 将两个集合合并
		 * @param a
		 * @param b
		 */
		public void union(V a, V b) {
			if (elementMap.containsKey(a) && elementMap.containsKey(b)) { // 检查是否注册
				Element<V> aF = findHead(elementMap.get(a)); // 找代表元素
				Element<V> bF = findHead(elementMap.get(b));
				if (aF != bF) {
					Element<V> big = sizeMap.get(aF) >= sizeMap.get(bF) ? aF : bF;
					Element<V> small = big == aF ? bF : aF;
					fatherMap.put(small, big);
					sizeMap.put(big, sizeMap.get(aF) + sizeMap.get(bF));
					sizeMap.remove(small); // small再也不是代表节点了
				}
			}
		}

	}

}
